import numpy as np
import scipy.io as sio
import matplotlib.pyplot as plt

data1 = sio.loadmat('ex7data2.mat')
X = data1['X']
# （300，2）

plt.scatter(X[:, 0], X[:, 1])
plt.show()

# 1、获取每个样本所属的类别
def find_centroids(X, centros):
    idx = []  # X中每个样本最终所属的类别

    for i in range(len(X)):  # 300次循环
        # print('X[i]: ', X[i])
        # print('centros: ', centros)
        # print('X[i] - centros', X[i] - centros)
        dist = np.linalg.norm((X[i] - centros), axis=1)  # 距离
        id_i = np.argmin(dist)  # 找到最小索引
        idx.append(id_i)

    return np.array(idx)


centros = np.array([[3, 3], [6, 2], [8, 5]])
idx = find_centroids(X, centros)
print(idx)
# [0 2 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
#  0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
#  0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 1 0
#  0 0 1 0 0 0 0 0 0 0 1 0 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 1 0 0 0 0 1 0 0
#  0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 1 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0
#  0 0 1 0 0 0 0 0 0 0 0 1 0 0 0 1 1 1 1 1 1 1 1 1 1 2 1 1 1 1 1 1 1 2 1 1 1
#  1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 1 0 1 2 1 2 1 1 1 1 1 1 1 1 1 1 1 1 1
#  1 1 0 1 1 1 0 1 1 1 1 1 1 1 0 1 1 1 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 2
#  1 1 1 0]

# 2、计算聚类中心点
def compute_centros(X, idx, k):
    centros = []  # 聚类中心点

    for i in range(k):
        centros_i = np.mean(X[idx == i], axis=0)
        centros.append(centros_i)

    return np.array(centros)

print(compute_centros(X, idx, k=3))
# [[2.42830111 3.15792418]
#  [5.81350331 2.63365645]
#  [7.11938687 3.6166844 ]]


# 3、运行kmeans，重复执行1和2
def run_kmeans(X, centros, iters):
    k = len(centros)  # 类别数
    centros_all = []  # 每次迭代过程中的聚类中心点
    centros_all.append(centros)
    centros_i = centros

    for i in range(iters):
        idx = find_centroids(X, centros_i)  # 获取样本所属的类别
        centros_i = compute_centros(X, idx, k)  # 计算新的聚类中心点
        centros_all.append(centros_i)

    return idx, np.array(centros_all)

# 绘制数据集和聚类中心的移动轨迹
def plot_data(X, centros_all, idx):
    plt.figure()
    plt.scatter(X[:, 0], X[:, 1], c=idx, cmap='rainbow')
    plt.plot(centros_all[:, :, 0], centros_all[:, :, 1], 'kx--')
    plt.show()

idx, centros_all = run_kmeans(X, centros, iters=10)
plot_data(X, centros_all, idx)


# 观察初始聚类点的位置对聚类效果的影响
def init_centros(X, k):
    index = np.random.choice(len(X), k)
    return X[index]

print(init_centros(X, k=3))
# [[3.27844295 1.75043926]
#  [1.40260822 1.08726536]
#  [5.78769095 3.29255127]]

for i in range(4):
    idx, centros_all = run_kmeans(X, init_centros(X, k=3), iters=10)
    plot_data(X, centros_all, idx)


